package com.rong.boot.jdbc.datasource;

import cn.runsa.arch.boot.Context;
import cn.runsa.boot.extension.web.ContextRestEvent;
import cn.runsa.boot.web.pattern.SessionContextHolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.ApplicationListener;
import org.springframework.jdbc.datasource.ConnectionProxy;
import org.springframework.jdbc.datasource.DriverManagerDataSource;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Properties;

/**
 * 动态数据源，从线程上下文中获取数据库名称 platform
 */
public class MultiTenantDataSource  extends DriverManagerDataSource implements ApplicationListener<ContextRestEvent> {
    private static final Logger log = LoggerFactory.getLogger(MultiTenantDataSource.class);
    private static final ThreadLocal<MultiTenantDataSource.ConnectionInvocationHandler> CONNECTION = new ThreadLocal();
    public static final int CONNECTION_MAX_LIVE_TIME = 120000;

    public MultiTenantDataSource() {
    }

    protected Connection getConnectionFromDriver(Properties props) throws SQLException {
        Context context = SessionContextHolder.getContext();
        if (context != null && context.getPlatform() != null) {
            String platform = context.getPlatform();
            MultiTenantDataSource.ConnectionInvocationHandler handler = (MultiTenantDataSource.ConnectionInvocationHandler)CONNECTION.get();
            if (handler != null) {
                if (platform.equals(handler.getPlatform())) {
                    if (System.currentTimeMillis() - handler.startTime <= 120000L) {
                        return (Connection) Proxy.newProxyInstance(ConnectionProxy.class.getClassLoader(), new Class[]{ConnectionProxy.class}, handler);
                    }

                    resetConnection();
                } else {
                    resetConnection();
                }
            }

            String url = this.getUrl() + "_" + context.getPlatform() + "?useUnicode=true&characterEncoding=UTF-8";
            if (log.isTraceEnabled()) {
                log.trace("establish db connection of {}", url);
            }

            Connection conn = super.getConnectionFromDriverManager(url, props);
            handler = new MultiTenantDataSource.ConnectionInvocationHandler(conn, null, platform);
            CONNECTION.set(handler);
            return (Connection)Proxy.newProxyInstance(ConnectionProxy.class.getClassLoader(), new Class[]{ConnectionProxy.class}, handler);
        } else {
            log.error(">> 配置数据源异常，未发现platform, 线程context="+context);
            throw new SQLException();
        }
    }

    public static void resetConnection() {
        MultiTenantDataSource.ConnectionInvocationHandler handler = (MultiTenantDataSource.ConnectionInvocationHandler)CONNECTION.get();
        if (handler != null) {
            if (log.isTraceEnabled()) {
                log.trace("release db connection,with tenant={}", handler.getTenant());
            }

            try {
                handler.target.close();
            } catch (Exception var2) {
                log.warn("close connection error", var2);
            }

            CONNECTION.remove();
        }

    }

    public void onApplicationEvent(ContextRestEvent event) {
        resetConnection();
    }

    static class ConnectionInvocationHandler implements InvocationHandler {
        private final Connection target;
        private final Integer tenant;
        private final String platform;
        private final long startTime;

        public ConnectionInvocationHandler(Connection target, Integer tenant, String platform) {
            this.target = target;
            this.tenant = tenant;
            this.platform = platform;
            this.startTime = System.currentTimeMillis();
        }

        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            if (method.getName().equals("equals")) {
                return proxy == args[0];
            } else if (method.getName().equals("hashCode")) {
                return System.identityHashCode(proxy);
            } else {
                if (method.getName().equals("unwrap")) {
                    if (((Class)args[0]).isInstance(proxy)) {
                        return proxy;
                    }
                } else if (method.getName().equals("isWrapperFor")) {
                    if (((Class)args[0]).isInstance(proxy)) {
                        return true;
                    }
                } else {
                    if (method.getName().equals("close")) {
                        return null;
                    }

                    if (method.getName().equals("isClosed")) {
                        return false;
                    }

                    if (method.getName().equals("getTargetConnection")) {
                        MultiTenantDataSource.log.debug("=====getTargetConnection");
                        return this.target;
                    }
                }

                try {
                    return method.invoke(this.target, args);
                } catch (InvocationTargetException var5) {
                    throw var5.getTargetException();
                }
            }
        }

        public Connection getTarget() {
            return this.target;
        }

        public Integer getTenant() {
            return this.tenant;
        }

        public long getStartTime() {
            return this.startTime;
        }

        public String getPlatform() {
            return this.platform;
        }
    }
}
